package two.chapter_4.c4_3;

import two.chapter_1.c1_3.Bag;
import two.chapter_1.c1_3.LinkedQueue;
import two.chapter_1.c1_3.Queue;
import two.chapter_2.c2_4.IndexMinPQ;

public class PrimMST extends MST {
    private Edge[]edgeTo;//距离树最近的边
    private double[] distTo;//distTo[w]=edgeTo[w].weight
    private boolean[] marked;//如果v在树中则为true
    private IndexMinPQ<Double> pq;
    public PrimMST(EdgeWeightedGraph g) {
        super(g);
        edgeTo=new Edge[g.v()];
        distTo=new double[g.v()];
        marked=new boolean[g.v()];
        pq=new IndexMinPQ<>(g.v());

        for (int v=0;v<g.v();v++){
            distTo[v]=Double.POSITIVE_INFINITY;
        }
        distTo[0]=0.0;
        pq.insert(0,0.0);
        while (!pq.isEmpty()){
            visit(g,pq.delMin());
        }
    }

    private void visit(EdgeWeightedGraph g, int v) {
        marked[v]=true;
        for (Edge e:g.adj(v)){
            int w=e.other(v);
            if (marked[w]) continue;
            if (e.weight()<distTo[w]){
                edgeTo[w]=e;
                distTo[w]=e.weight();
                if (pq.contains(w)){
                    pq.changeKey(w,distTo[w]);
                }else {
                    pq.insert(w,distTo[w]);
                }
            }
        }
    }

    @Override
    public Iterable<Edge> edges() {
        Queue<Edge> mst=new LinkedQueue<>();
        for (int v=0;v<edgeTo.length;v++){
            Edge e=edgeTo[v];
            if (e!=null){
                mst.enqueue(e);
            }
        }
        return mst;
    }


    @Override
    public double weight() {
        double weight = 0.0;
        for (Edge e : edges())
            weight += e.weight();
        return weight;
    }
}
